package org.thinkin.springboot.compiler;

import cn.hutool.core.io.IoUtil;
import lombok.SneakyThrows;

import javax.tools.JavaFileObject;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

public class ClassLoaderImpl extends ClassLoader{
    private final Map<String, JavaFileObject> classes = new HashMap<String, JavaFileObject>();
    private final String baseDir;
    public ClassLoaderImpl(final ClassLoader parentClassLoader,String baseDir) {
        super(parentClassLoader);
        this.baseDir = baseDir;
    }

    Collection<JavaFileObject> files() {
        return Collections.unmodifiableCollection(classes.values());
    }

    @SneakyThrows
    @Override
    protected Class<?> findClass(final String qualifiedClassName)  {
        JavaFileObject file = classes.get(qualifiedClassName);
        if (file != null) {
            byte[] bytes = ((JavaFileObjectImpl) file).getByteCode();
            return defineClass(qualifiedClassName, bytes, 0, bytes.length);
        }
        else {
            file = new JavaFileObjectImpl(qualifiedClassName, JavaFileObject.Kind.CLASS);
            OutputStream out = file.openOutputStream();
            byte[] bytes1 = IoUtil.readBytes(Files.newInputStream(
                    new File(baseDir, qualifiedClassName + ClassUtils.CLASS_EXTENSION).toPath()));
            out.write(bytes1);
            this.add(qualifiedClassName,file);
            byte[] bytes = ((JavaFileObjectImpl) file).getByteCode();
            return defineClass(qualifiedClassName, bytes, 0, bytes.length);
        }
    }

    void add(final String qualifiedClassName, final JavaFileObject javaFile) {
        classes.put(qualifiedClassName, javaFile);
    }

    @Override
    protected synchronized Class<?> loadClass(final String name, final boolean resolve) throws ClassNotFoundException {
        return super.loadClass(name, resolve);
    }

    @Override
    public InputStream getResourceAsStream(final String name) {
        if (name.endsWith(ClassUtils.CLASS_EXTENSION)) {
            String qualifiedClassName = name.substring(0, name.length() - ClassUtils.CLASS_EXTENSION.length())
                    .replace('/', '.');
            JavaFileObjectImpl file = (JavaFileObjectImpl) classes.get(qualifiedClassName);
            if (file != null) {
                return new ByteArrayInputStream(file.getByteCode());
            }
        }
        return super.getResourceAsStream(name);
    }
}
